#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：cls_template 
@File    ：predict.py
@Author  ：ChenmingSong
@Date    ：2022/1/5 16:23 
@Description：用来推理数据集
'''
import torch
# from train_resnet import SelfNet
from train import SELFMODEL
import os
import os.path as osp
import shutil
from PIL import Image
from utils.transforms import get_torch_transforms

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

model_path = "../checkpoints/resnet50_pretrained_224/resnet50_9epochs_accuracy0.92633_weights.pth"  # todo  模型路径
classes_names = ['一品红', '万寿菊', '乌头', '仙客来', '六出花', '凌霄', '凤梨', '刺芹', '勋章菊', '卡特兰', '卷丹', '叶子花', '向日葵', '唐菖蒲', '嘉兰', '大丽花', '大滨菊', '天人菊', '天竺葵', '天蓝绣球', '姜荷花', '射干', '山桃草', '山茶花', '帝王花', '康乃馨', '德国鸢尾', '报春花', '旱金莲', '曼陀罗', '月季', '朱槿', '杜鹃花', '松果菊', '桔梗', '款冬', '毛地黄', '毛茛', '水塔花', '沙漠玫瑰', '油点草', '洋桔梗', '牵牛花', '玫瑰', '番红花', '睡莲', '石竹', '石蒜', '硬叶兜兰', '碧冬茄', '秋水仙', '秋英', '糖芥', '红蕉', '罂粟', '美丽月见草', '美人蕉', '美国薄荷', '耧斗花', '芍药', '芙蓉葵', '花毛茛', '花烛', '花菱草', '荷花', '菊芋', '葡萄风信子', '蒲公英', '蓝刺头', '蓝盆花', '蓝目菊', '蓝花矢车菊', '蓝花草', '蓟', '蝴蝶兰', '西番莲', '角堇', '贝母', '郁金香', '野棉花', '金盏花', '金脉爵床', '金鱼草', '铁筷子', '铁线莲', '银灰旋花', '银莲花', '非洲菊', '风铃草', '香豌豆', '马蹄莲', '鸡蛋花', '鹤望兰', '鹤草', '黄水仙', '黄菖蒲', '黑心金光菊', '龙胆']  # todo 类名
img_size = 224  # todo 图片大小
model_name = "resnet50"  # todo 模型名称
num_classes = len(classes_names)  # todo 类别数目


def predict_batch(model_path, target_dir, save_dir):
    data_transforms = get_torch_transforms(img_size=img_size)
    valid_transforms = data_transforms['val']
    # 加载网络
    model = SELFMODEL(model_name=model_name, out_features=num_classes, pretrained=False)
    # model = nn.DataParallel(model)
    weights = torch.load(model_path)
    model.load_state_dict(weights)
    model.eval()
    model.to(device)
    # 读取图片
    image_names = os.listdir(target_dir)
    for i, image_name in enumerate(image_names):
        image_path = osp.join(target_dir, image_name)
        img = Image.open(image_path)
        img = valid_transforms(img)
        img = img.unsqueeze(0)
        img = img.to(device)
        output = model(img)
        label_id = torch.argmax(output).item()
        predict_name = classes_names[label_id]
        save_path = osp.join(save_dir, predict_name)
        if not osp.isdir(save_path):
            os.makedirs(save_path)
        shutil.copy(image_path, save_path)
        print(f"{i + 1}: {image_name} result {predict_name}")


def predict_single(model_path, image_path):
    data_transforms = get_torch_transforms(img_size=img_size)
    # train_transforms = data_transforms['train']
    valid_transforms = data_transforms['val']
    # 加载网络
    model = SELFMODEL(model_name=model_name, out_features=num_classes, pretrained=False)
    # model = nn.DataParallel(model)
    weights = torch.load(model_path)
    model.load_state_dict(weights)
    model.eval()
    model.to(device)

    # 读取图片
    img = Image.open(image_path)
    img = valid_transforms(img)
    img = img.unsqueeze(0)
    img = img.to(device)
    output = model(img)
    label_id = torch.argmax(output).item()
    predict_name = classes_names[label_id]
    print(f"{image_path}'s result is {predict_name}")


if __name__ == '__main__':
    # 批量预测函数
    # predict_batch(model_path=model_path,
    #               target_dir="D:/upppppppppp/cls/cls_torch_tem/images/test_imgs/mini",
    #               save_dir="D:/upppppppppp/cls/cls_torch_tem/images/test_imgs/mini_result")
    # 单张图片预测函数
    predict_single(model_path=model_path, image_path="resources/images/test_imgs/506659320_6fac46551e.jpg")

